#!/usr/bin/env python
# coding: utf-8

# In[ ]:


import random
import matplotlib.pyplot as plt


# In[43]:


# 种群初始基因型
AA = 30000
Aa = 60000
aa = 10000


# In[ ]:


#模拟分离定律
def segregation(n):
    results = []
    for _ in range(n):
        r = random.randint(0,3)
        if r==0:    results.append(0)
        if r in [1,2]:    results.append(1)
        if r==3:    results.append(2)
    return results


# In[ ]:


# 根据亲本基因型确定子代基因型
def GenerateChild(parents):
    ChildNum = random.randint(1,2) #子代个数
    parents.sort()
    if parents[0]==0 and parents[1]==0: # AA*AA
        return [0]*ChildNum
    elif parents[0]==1 and parents[1]==1: # Aa*Aa
        return segregation(ChildNum)
    elif parents[0]==2 and parents[1]==2: # aa*aa
        return [2]*ChildNum
    elif parents[0]==0 and parents[1]==1: # AA*Aa
        return [random.choice([0,1]) for _ in range(ChildNum)]
    elif parents[0]==0 and parents[1]==2: # AA*aa
        return [1]*ChildNum
    elif parents[0]==1 and parents[1]==2: # Aa*aa
        return [random.choice([1,2]) for _ in range(ChildNum)]
    


# In[13]:


#自由交配
def FreeMating():
    global GenePool
    children = []
    for _ in range(int(len(GenePool)*0.67)): #模拟自由交配数为个体总数的0.67倍，以控制子代数量
        children.extend(GenerateChild(random.choices(GenePool,k=2)))
    GenePool = list(children)


# In[44]:


GenePool = [0 for _ in range(AA)] + [1 for _ in range(Aa)] + [2 for _ in range(aa)] #重置基因库

HisFreq = [] # 记录每代基因型数量
for _ in range(100):
    HisFreq.append((GenePool.count(0),GenePool.count(1),GenePool.count(2)))
    FreeMating()
    print(str(len(GenePool))+',n='+str(_))
print("AA\tAa\taa")
for t in HisFreq:
    print(str(round(t[0]/sum(t),3))+'\t'+str(round(t[1]/sum(t),3))+'\t'+str(round(t[2]/sum(t),3)))


# In[45]:


#绘图，基因型频率
plt.plot(range(len(HisFreq)),[(t[0]/sum(t)) for t in HisFreq],'r',label="AA")
plt.plot(range(len(HisFreq)),[(t[1]/sum(t)) for t in HisFreq],'g',label="Aa")
plt.plot(range(len(HisFreq)),[(t[2]/sum(t)) for t in HisFreq],'b',label="aa")
plt.xlabel("Generation")
plt.ylabel("Genotype frequency")
plt.legend(loc=1)
plt.savefig('./a0000.jpg',dpi=300)
plt.show()


# In[46]:


#绘图，基因频率

HisGenotype = [(int(x+0.5*y),int(z+0.5*y)) for x,y,z in HisFreq] # 根据基因型数量计算基因数量
plt.plot(range(len(HisGenotype)),[(t[0]/sum(t)) for t in HisGenotype],'b',label="A")
plt.plot(range(len(HisGenotype)),[(t[1]/sum(t)) for t in HisGenotype],'g',label="a")
plt.xlabel("Generation")
plt.ylabel("Gene frequency")
plt.legend(loc=1)
plt.savefig('./b0000.jpg',dpi=300)
plt.show()

